import math
from ucimlrepo import fetch_ucirepo

# 获取数据集
wine = fetch_ucirepo(id=109)

# 数据（作为 pandas 数据帧）
X = wine.data.features
y = wine.data.targets

# 合并特征和目标标签
data = X.join(y)

# 从元数据中提取属性名称和标签
attributes = list(X.columns) + ['标签']

class Node:
    def __init__(self, attribute=None, label=None, branches=None):
        self.attribute = attribute  # 分裂属性
        self.label = label  # 叶节点标签
        self.branches = branches or {}  # 分支，存储子节点

def entropy(data):
    """计算数据集的信息熵"""
    total_instances = len(data)
    if total_instances == 0:
        return 0

    label_counts = data.iloc[:, -1].value_counts()
    entropy_value = 0
    for count in label_counts:
        probability = count / total_instances
        entropy_value -= probability * math.log2(probability)

    return entropy_value

def gain(data, attribute):
    """计算特定属性的信息增益"""
    total_instances = len(data)
    if total_instances == 0:
        return 0

    attribute_values = data[attribute].unique()
    subset_entropy = 0

    for value in attribute_values:
        subset = data[data[attribute] == value]
        subset_probability = len(subset) / total_instances
        subset_entropy += subset_probability * entropy(subset)

    information_gain = entropy(data) - subset_entropy
    return information_gain

def get_best_attribute(data, attributes):
    """选择信息增益最大的属性"""
    information_gains = {attribute: gain(data, attribute) for attribute in attributes[:-1]}
    best_attribute = max(information_gains, key=information_gains.get)
    return best_attribute

def id3(data, attributes):
    """使用 ID3 算法构建决策树"""
    labels = data.iloc[:, -1].tolist()

    # 如果所有实例属于同一类别，则返回叶节点
    if len(set(labels)) == 1:
        return Node(label=labels[0])

    # 如果属性集为空，则返回叶节点，选择类别标签为数据中出现最多的类别
    if len(attributes) == 1:
        majority_label = max(set(labels), key=labels.count)
        return Node(label=majority_label)

    # 选择信息增益最大的属性
    best_attribute = get_best_attribute(data, attributes)

    # 构建树
    branches = {}
    for value in data[best_attribute].unique():
        default_label = max(set(labels), key=labels.count)
        subset = data[data[best_attribute] == value]
        if subset.empty:
            branches[value] = Node(label=default_label)
        else:
            new_attributes = [attr for attr in attributes if attr != best_attribute]
            branches[value] = id3(subset, new_attributes)

    return Node(attribute=best_attribute, branches=branches)

def print_tree(node, depth=0):
    """以可读格式打印决策树"""
    if node.label is not None:
        print("  " * depth + "标签: " + str(node.label))
    else:
        print("  " * depth + "属性: " + str(node.attribute))
        for value, branch in node.branches.items():
            print("  " * (depth + 1) + "值: " + str(value))
            print_tree(branch, depth + 2)

# 构建决策树
decision_tree = id3(data, attributes)

# 打印决策树
print_tree(decision_tree)